import os
from PyQt5.QtWidgets import QMenu,QWidget,QListWidget
from PyQt5.Qt import QPoint,Qt,QPainter,QPointF,pyqtSignal,QThread
from PyQt5.QtGui import QImage, QPixmap,QPen,QColor,QCursor,QPalette,QBrush
from PyUIC.QpaintWidget import Ui_Form
import cv2
# from QWidgets.QLabelDialog import QLabelDialog

def distance(p,points):
    if (p.x()>points[0] and p.x()<points[2]) and (p.y()>points[1] and p.y()<points[3]):
        return True
    else:
        return False

class PaintThread(QThread):
    trigger = pyqtSignal(str)
    def __int__(self,image):
        # 初始化函数
        super(PaintThread, self).__init__()
        self.image = image

    def run(self):
        self.h_vertex = None
        self.h_shape = None
        image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        x = image.shape[1]
        y = image.shape[0]

        self.zoomscale = 1
        frame = QImage(image, image.shape[1], image.shape[0], image.shape[1] * 3, QImage.Format_RGB888)
        self.img = QPixmap.fromImage(frame)
        self.scale = min(self.height() / self.img.height(), self.width() / self.img.width())
        # self.img = QPixmap(image_path)
        self.scaled_img = self.img.scaled(self.scale,Qt.IgnoreAspectRatio)

class ImageBox(QWidget):
    enddrawing = pyqtSignal()
    def __init__(self):
        super(ImageBox,self).__init__()
        self._painter = QPainter()
        self.img = None
        self.scaled_img = None
        self.point = QPoint(0,0)
        self.start_pos = None
        self.end_pos = None
        self.left_click = False
        self.scale =1
        self.rect_list = []
        self.rect=None
        self.roi_box = None
        self.is_show_roi = True
        self.check_state = False
        self.h_shape = None
        self.h_vertex = None

        # self.size = (500,500)
    # def init_ui(self):
    def set_image(self,image):
        # image = cv2.imread(image_path)
        self.h_vertex=None
        self.h_shape=None
        image= cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        x = image.shape[1]
        y = image.shape[0]

        self.zoomscale=1
        frame = QImage(image,image.shape[1], image.shape[0], image.shape[1] * 3,QImage.Format_RGB888)
        self.img = QPixmap.fromImage(frame)
        self.scale= min(self.height() / self.img.height(), self.width() / self.img.width())
        # self.img = QPixmap(image_path)
        self.scaled_img = self.img.scaled(self.scale,Qt.IgnoreAspectRatio)
        self.repaint()

    def addRects(self,rects,labels):
        if type(rects).__name__ == 'list':
            self.rect_list = rects
            self.labels = labels
            self.repaint()
        else:
            raise TypeError('rects is not a list')
    def addRect(self,x):
        # self.rect = (int(x[0]),int(x[1]),int(x[2])-int(x[0]),int(x[3])-int(x[1]))
        self.rect_list.append(x)

    def addRoiBox(self,roi_box):
        self.roi_box = (int(roi_box[2]),int(roi_box[0]),int(roi_box[3])-int(roi_box[2]),int(roi_box[1])-int(roi_box[0]))

    def setRect(self,x):

        # print('Rect',x)
        self.rect = (int(x[0]),int(x[1]),int(x[2])-int(x[0]),int(x[3])-int(x[1]))

        # print(self.rect)
        self.repaint()

    def nearest_vertex(self,point, epsilon):
        for i, p in enumerate(self.rect_list):
            # print(p)
            if distance(point,p):
                return i
        return None

    def paintEvent(self,e):
        if self.img:
            painter = self._painter
            painter.begin(self)
            painter.setRenderHint(QPainter.Antialiasing)
            painter.setRenderHint(QPainter.HighQualityAntialiasing)
            painter.setRenderHint(QPainter.SmoothPixmapTransform)

            # print('height width',min(self.height(),self.width()))

            # self.scale_x,self.scale_y = self.height()/self.img.height(),self.width()/self.img.width()
            # print(self.scale_x,self.scale_y)

            painter.scale(self.scale,self.scale)
            painter.translate(self.offset_to_center())

            painter.drawPixmap(0,0,self.img)
            label_font_size = 8


            if len(self.rect_list)>0:
                self.drawRectList(painter)
            if self.rect:
                self.drawRect(painter)
            if self.roi_box and self.is_show_roi:
                pen = QPen(Qt.yellow, 4)
                painter.setPen(pen)
                painter.drawRect(*self.roi_box)
            if self.h_shape:
                color = QColor(0, 128, 255, 155)
                pen = QPen(color, 4)
                painter.setPen(pen)
                rect = (int(self.h_shape[0]), int(self.h_shape[1]), int(self.h_shape[2]) - int(self.h_shape[0]), int(self.h_shape[3]) - int(self.h_shape[1]))
                # painter.drawRect(*rect)
                painter.fillRect(*rect,color)
            painter.end()

    def transform_pos(self, point):
        """Convert from widget-logical coordinates to painter-logical coordinates."""
        return point / self.scale - self.offset_to_center()

    def offset_to_center(self):
        s = self.scale
        area = super(ImageBox, self).size()
        w, h = self.img.width() * s, self.img.height() * s
        aw, ah = area.width(), area.height()
        x = (aw - w) / (2 * s) if aw > w else 0
        y = (ah - h) / (2 * s) if ah > h else 0
        return QPointF(x, y)


    def mouseMoveEvent(self,e):
        if self.img:
            pos = self.transform_pos(e.pos())
            window = self.parent().window()
            if self.check_state:
                start_x, start_y = self.rect[0:2]
                # s = self.scale
                # area = super(ImageBox, self).size()
                # w, h = self.img.width() * s, self.img.height() * s
                # aw, ah = area.width(), area.height()
                # x = (aw - w) / (2 * s) if aw > w else 0
                # y = (ah - h) / (2 * s) if ah > h else 0
                self.rect = (start_x, start_y, pos.x()- start_x, pos.y() - start_y)
                self.update()
            # else:
            #     index = self.nearest_vertex(point=pos, epsilon=11.0)
            #     print("index", index)
            #     if index is not None:
            #         self.h_vertex, self.h_shape = index, self.rect_list[index]
            #     self.repaint()

        # def keyPressEvent(self,e):
        #     if e.button() = Qt.

    def mousePressEvent(self, e):
        if self.img:
            pos = self.transform_pos(e.pos())
            if e.button() == Qt.LeftButton:
                if self.check_state:
                    self.rect = (pos.x(),pos.y(), 0, 0)
                else:

                    index = self.nearest_vertex(point=pos, epsilon=11.0)
                    # print("index", index)
                    if index is not None:
                        self.h_vertex, self.h_shape = index, self.rect_list[index]
                    self.repaint()

    def mouseReleaseEvent(self,e):
        if self.img:
            if e.button() == Qt.LeftButton:
                if self.check_state:

                    self.enddrawing.emit()
                    self.rect = None

                    self.update()


    def wheelEvent(self,e):
        pass
        # if self.img is None:
        #     return
        # angle = e.angleDelta()/8
        # angleX = angle.x()
        # angleY = angle.y()
        #
        # if angleY>0:
        #     # print("鼠标中键上滚")
        #     self.scale *= 1.1
        #     # self.scale_y *= 1.1
        #     # self.scaled_img = self.img.scaled(self.scaled_img.width()+10,self.scaled_img.height()+10)
        #     # newWidth = e.x() - (self.scaled_img.width()*(e.x()-self.point.x()))/(self.scaled_img.width()-10)
        #     # newHeihght = e.y() - (self.scaled_img.height()*(e.y()-self.point.y()))/(self.scaled_img.height()-10)
        #     # self.singleOffset = QPoint(newWidth,newHeihght)
        #     # start_x, start_y ,end_x,end_y = self.rect[:]
        #     # start_x *= newWidth / self.scaled_img.width()
        #     # start_y *= newHeihght / self.scaled_img.height()
        #     # end_x *= newWidth / self.scaled_img.width()
        #     # end_y *= newHeihght / self.scaled_img.height()
        #     # self.rect = (start_x, start_y ,end_x,end_y)
        #     self.repaint()
        # else:
        #     # print("鼠标中键上滚")
        #     self.scale /=1.1
        #     # self.scale_y /= 1.1
        #     # self.scaled_img = self.img.scaled(self.scaled_img.width() *1.1, self.scaled_img.height()*1.1)
        #     # newWidth = e.x() - (self.scaled_img.width() * (e.x() - self.point.x())) / (
        #     #             self.scaled_img.width() + 10)
        #     # newHeihght = e.y() - (self.scaled_img.height() * (e.y() - self.point.y())) / (
        #     #             self.scaled_img.height() + 10)
        #     # self.point = QPoint(newWidth, newHeihght)
        #     # start_x, start_y, end_x, end_y = self.rect[:]
        #     # start_x *= newWidth / self.scaled_img.width()
        #     # start_y *= newHeihght / self.scaled_img.height()
        #     # end_x *= newWidth / self.scaled_img.width()
        #     # end_y *= newHeihght / self.scaled_img.height()
        #     # self.rect = (start_x, start_y, end_x, end_y)
        #     self.repaint()

    def drawRect(self,qp):
        #创建红色 宽度为2个像素的画笔
        pen = QPen(Qt.green,2)
        qp.setPen(pen)
        qp.drawRect(*self.rect)

    def drawRectList(self,qp):
        pen = QPen(Qt.red,4)
        qp.setPen(pen)
        for x in self.rect_list:
            rect = (int(x[0]),int(x[1]),int(x[2])-int(x[0]),int(x[3])-int(x[1]))
            qp.drawRect(*rect)





class QPaintWidgetView(QWidget,Ui_Form):
    edit_label = pyqtSignal(bool)
    del_shape = pyqtSignal(int)
    def __init__(self,parent=None):
        super(QPaintWidgetView,self).__init__(parent)
        self.setupUi(self)
        # pal = QPalette()
        # pal.setBrush(QPalette.Background, QColor(150,150,150,80))
        # self.setPalette(pal)
        # self.setAutoFillBackground(True)
        # self.setStyleSheet("border:5px solid #ffffff;#background:rgb(150,150,150,80)")
        # self.AoiBox.toggled.connect(self.ShowAOIBox)
        self.label_list = QListWidget()
        self.train_txt_path = 'train/train.txt'
        self.img_paths=[]
        self.label_txt_folder = './label'
        # self.setFixedSize(800,800)
        self.setAutoFillBackground(True)
        # print('height',self.height(),self.width())
        self._size = min(self.height(),self.width())
        # print('_size', self._size)
        self.box = ImageBox()
        # self.box.resize(self._size,self._size)
        # self.SaveButton.clicked.connect(self.SaveSlot)
        # self.CreateRectButton.clicked.connect(self.ChangeState)
        self.box.enddrawing.connect(self.EndDrawing)
        self.PaintLayout.addWidget(self.box)
        # self.label_dialog = QLabelDialog(parent=self, list_item=[])

        # self.setContextMenuPolicy(Qt.CustomContextMenu)
        # 创建QMenu
        # self.customContextMenuRequested.connect(self.showMenu)

        # self.contextMenu = QMenu()
        # self.CreateShape = self.contextMenu.addAction("创建新形状")
        # self.DeleteShape = self.contextMenu.addAction("删除选中形状")
        # self.CreateShapeShortcut = QShortcut(QKeySequence("Ctrl+N"),self)
        # self.DeleteShapeShortcut = QShortcut(QKeySequence("Ctrl+D"),self)
        # self.CreateShape.triggered.connect(self.ChangeState)
        # self.DeleteShape.triggered.connect(self.DeleteShapeSlot)
        # self.CreateShapeShortcut.activated.connect(self.ChangeState)
        # self.DeleteShapeShortcut.activated.connect(self.DeleteShapeSlot)

    def open_image(self,image):
        # self.img_path = img_path
        # if self.img_path in self.img_paths:
        #     print("图像已存在")
        # else:
        #     self.img_paths.append(self.img_path)
        self.box.set_image(image)

    def setRect(self,x):
        self.box.setRect(x)

    def getNewRect(self):
        return self.box.rect

    def ChangeState(self):
        self.CreateRectButton.setEnabled(False)
        self.box.check_state=True
        self.setCursor(Qt.CrossCursor)

    def ShowAOIBox(self):
        if self.AoiBox.isChecked():
            self.box.is_show_roi = True
        else:
            self.box.is_show_roi = False
        self.box.repaint()
    def DeleteShapeSlot(self):
        shape = self.box.h_shape
        index = self.box.h_vertex
        total = len(self.box.rect_list)
        if shape is None:
            return
        print(shape,'\n',total-1-index)
        self.box.h_shape = None
        self.box.h_vertex = None
        self.box.rect=None
        self.del_shape.emit(total-1-index)



    def current_item(self):
        items = self.label_list.selectedItems()
        if items:
            return items[0]
        return None

    def EndDrawing(self):
        self.CreateRectButton.setEnabled(True)
        self.box.check_state = False
        self.setCursor(Qt.ArrowCursor)
        self.edit_label.emit(True)

    def SaveSlot(self):
        if len(self.box.rect_list)>0 and self.box.img is not None:
            self.save_train_txt(self.img_path)

            rects = self.box.rect_list
            labels = self.box.labels
            print(rects)
            self.save_label_txt(rects,labels)

    def save_train_txt(self,img_path):
        print('Saving train_txt ... ')
        if self.train_txt_path is not None:
            if img_path in self.img_paths:
                print("%s 已保存" % img_path)
            else:
                self.img_paths.append(self.img_path)
                with open(self.train_txt_path,"a",encoding='utf-8') as file:
                    file.write(self.img_path+'\n')
        return

    def save_label_txt(self,rects,labels):
        print('Saving label_txt ... ')
        if not os.path.exists(self.label_txt_folder):
            os.makedirs(self.label_txt_folder)
        img_name = os.path.basename(self.img_path)
        label_name = img_name.replace('jpg','txt')
        label_name = os.path.join(self.label_txt_folder,label_name)
        # print(label_name)
        if label_name is not None:
            with open(label_name,"w",encoding='utf-8') as file:
               for i,rect in enumerate(rects):
                   # print(rect)
                   b1,b2,b3,b4 = int(rect[0]),int(rect[2]),int(rect[1]),int(rect[3])
                   b = (b1,b2,b3,b4)
                   bb = self.convert((2048,2048),b)
                   # print(bb)
                   file.write(str(labels[i])+" "+ " ".join([str(a) for a in bb])+ '\n')
        return

    def convert(self,size,box):
        dw = 1./(size[0])
        dh = 1./(size[1])
        x = (box[0]+box[1])/2.0 -1
        y = (box[2]+box[3])/2.0 -1
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = x*dw
        w = w*dw
        y = y*dh
        h = h*dh
        return x,y,w,h

    def showMenu(self,pos):
        self.contextMenu.exec_(QCursor.pos())







